import argparse
from server.model import gentext,genimg,tts,removebg




if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Predict Tokens using `generate()` API for GLM-4 or SDXL model')
    parser.add_argument('--task', type=str,  default="text",
                        help='text,image, tts or remove')
    parser.add_argument('--input', type=str, default="请帮我生成一篇关于科技兴国的公众号文章",
                        help='Prompt to infer or text to synthesize or image to removebg')
    

    args = parser.parse_args()
    # 文本生成
    if args.task == "text":
        gentext(args.input)
    # 图片生成
    if args.model == "image":
        genimg(args.input)
    # 虚拟人语音转文字
    if args.model == "tts":
        tts(args.input)
    # 智能抠图
    if args.model == "remove":
        removebg(args.input)
